import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import least_squares
from scipy.integrate import quad


Data_list = ['100K_Si111.txt','300K_Si111.txt','400K_Si111.txt','500K_Si111.txt','600K_Si111.txt','680K_Si111.txt']

MT_list = ['300K_Si111_MT.txt','300K_Si111_MT.txt','400K_Si111_MT.txt','500K_Si111_MT.txt','600K_Si111_MT.txt','680K_Si111_MT.txt']

E_length = 600 + 1


i = 0
for fname in Data_list:
    Data = []
    count = 1
    with open(fname, 'r') as D:
        next(D)
        for line in D:
            if count%E_length != 1:
                Data.append([float(x) for x in line.split(',')])
            count+=1
    Data = np.array(Data)
    locals()['Data'+str(i)] = Data    
    i+=1

i = 0
for fname in MT_list:
    MT = []
    count = 1
    with open(fname, 'r') as D:
        next(D)
        for line in D:
            if count%E_length != 1:
                MT.append([float(x) for x in line.split(',')])
            count+=1
    MT = np.array(MT)
    locals()['MT'+str(i)] = MT   
    i+=1

R0 = np.max(Data0[3600:4200,1]-MT0[3600:4200,1]*0.75)
R1 = np.max(Data1[3600:4200,1]-MT1[3600:4200,1]*0.75)/R0
R2 = np.max(Data2[3600:4200,1]-MT2[3600:4200,1]*0.75)/R0
R3 = np.max(Data3[3600:4200,1]-MT3[3600:4200,1]*0.75)/R0
R4 = np.max(Data4[3600:4200,1]-MT4[3600:4200,1]*0.75)/R0
R5 = np.max(Data5[3600:4200,1]-MT5[3600:4200,1]*0.75)/R0

#2.995775,2.581035,1.592615,1.36199,1.2151675000000002,1.2682475000000002

plt.errorbar(Data0[3600:4200,0]*1000,Data0[3600:4200,1]-MT0[3600:4200,1]*0.75,Data0[3600:4200,2]-MT0[3600:4200,2]*0.75,color='black')
plt.errorbar(Data1[3600:4200,0]*1000,(Data1[3600:4200,1]-MT1[3600:4200,1]*0.75)+0.01,Data1[3600:4200,2]-MT1[3600:4200,2]*0.75,color='C0')
plt.errorbar(Data2[3600:4200,0]*1000,(Data2[3600:4200,1]-MT2[3600:4200,1]*0.75)+0.02,Data2[3600:4200,2]-MT2[3600:4200,2]*0.75,color='C2')
plt.errorbar(Data3[3600:4200,0]*1000,(Data3[3600:4200,1]-MT3[3600:4200,1]*0.75)*2+0.03,Data3[3600:4200,2]-MT3[3600:4200,2]*0.75,color='C1')
plt.errorbar(Data4[3600:4200,0]*1000,(Data4[3600:4200,1]-MT4[3600:4200,1]*0.75)*3+0.04,Data4[3600:4200,2]-MT4[3600:4200,2]*0.75,color='C3')
plt.errorbar(Data5[3600:4200,0]*1000,(Data5[3600:4200,1]-MT5[3600:4200,1]*0.75)*4+0.05,Data5[3600:4200,2]-MT5[3600:4200,2]*0.75,color='C4')

# norm to 1
#plt.errorbar(Data0[3600:4200,0]*1000,(Data0[3600:4200,1]-MT0[3600:4200,1]*0.75)/np.max(Data0[3600:4200,1]-MT0[3600:4200,1]*0.75),(Data0[3600:4200,2]-MT0[3600:4200,2]*0.75)/np.max(Data0[3600:4200,1]-MT0[3600:4200,1]*0.75))
#plt.errorbar(Data1[3600:4200,0]*1000,(Data1[3600:4200,1]-MT1[3600:4200,1]*0.75)/np.max(Data1[3600:4200,1]-MT1[3600:4200,1]*0.75),(Data1[3600:4200,2]-MT1[3600:4200,2]*0.75)/np.max(Data1[3600:4200,1]-MT1[3600:4200,1]*0.75))
#plt.errorbar(Data2[3600:4200,0]*1000,(Data2[3600:4200,1]-MT2[3600:4200,1]*0.75)/np.max(Data2[3600:4200,1]-MT2[3600:4200,1]*0.75),(Data2[3600:4200,2]-MT2[3600:4200,2]*0.75)/np.max(Data2[3600:4200,1]-MT2[3600:4200,1]*0.75))
#plt.errorbar(Data3[3600:4200,0]*1000,(Data3[3600:4200,1]-MT3[3600:4200,1]*0.75)/np.max(Data3[3600:4200,1]-MT3[3600:4200,1]*0.75),(Data3[3600:4200,2]-MT3[3600:4200,2]*0.75)/np.max(Data3[3600:4200,1]-MT3[3600:4200,1]*0.75))
#plt.errorbar(Data4[3600:4200,0]*1000,(Data4[3600:4200,1]-MT4[3600:4200,1]*0.75)/np.max(Data4[3600:4200,1]-MT4[3600:4200,1]*0.75),(Data4[3600:4200,2]-MT4[3600:4200,2]*0.75)/np.max(Data4[3600:4200,1]-MT4[3600:4200,1]*0.75))

plt.text(65,0.0015,r'100K',color='black',fontsize=20,ha='left',va='bottom')
plt.text(65,0.0125,r'300K',color='C0',fontsize=20,ha='left',va='bottom')
plt.text(65,0.024,r'400K',color='C2',fontsize=20,ha='left',va='bottom')
plt.text(65,0.040,r'500K',color='C1',fontsize=20,ha='left',va='bottom')
plt.text(65,0.052,r'600K',color='C3',fontsize=20,ha='left',va='bottom')
plt.text(65,0.066,r'680K',color='C4',fontsize=20,ha='left',va='bottom')

plt.text(98,0.042,r'$\times$2',color='C1',fontsize=12,ha='left',va='bottom')
plt.text(98,0.0535,r'$\times$3',color='C3',fontsize=12,ha='left',va='bottom')
plt.text(98,0.068,r'$\times$4',color='C4',fontsize=12,ha='left',va='bottom')

#plt.legend(['300K','400K','500K','600K','680K'],frameon=False,fontsize=16,ncol=3,loc='upper right')
plt.title(r'Q = 1.5$\pm0.1 \rm{\AA}^{-1}$',fontsize=20)
#plt.semilogy()
plt.xlim(-110,110)
plt.ylim(0,0.08)
plt.yticks([0,0.02,0.04,0.06,0.08])
plt.xlabel(r'Energy ($\mu$eV)',fontsize=20)
plt.ylabel('Intensity (a.u.)',fontsize=20)
plt.tick_params(which='both',direction='in',labelsize=20,pad=6)
plt.savefig('overplotIE_v3.eps',format='eps',bbox_inches='tight')
plt.show()
